# 本示例演示如何使用kan进行分类
import torch
import numpy as np
from   kan import KAN

import matplotlib.pyplot as plt
from   sklearn.datasets  import make_moons
from   sklearn.ensemble  import RandomForestClassifier


def create_dataset(train_samples=1000, test_samples=1000, shuffle=True, extend=True):
    dataset = {}
    train_input, train_label = make_moons(n_samples=train_samples, shuffle=shuffle, 
                                        noise=0.1, random_state=None)
    test_input,  test_label  = make_moons(n_samples=test_samples,  shuffle=shuffle, 
                                        noise=0.1, random_state=None)

    dataset['train_input']       = torch.from_numpy(train_input)
    dataset['test_input']        = torch.from_numpy(test_input)

    if extend:
        dataset['train_label']   = torch.from_numpy(train_label[:,None])
        dataset['test_label']    = torch.from_numpy(test_label[:,None])
    else:
        dataset['train_label']   = torch.from_numpy(train_label)
        dataset['test_label']    = torch.from_numpy(test_label)

    # if plot_train:
    #     X = dataset['train_input']
    #     y = dataset['train_label']
    #     plt.scatter(X[:,0], X[:,1], c=y[:,0])
    #     plt.show()
    
    return dataset


def random_forest(dataset_fn, n_estimators=1, max_depth=10, random_state=50, n_jobs=4):
    model = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth, 
                                   random_state=random_state, n_jobs=n_jobs)
    dataset = dataset_fn(extend=False)
    model.fit(dataset["train_input"], dataset["train_label"])

    accuracy = model.score(dataset["train_input"], dataset["train_label"])
    print("模型在训练数据集上的精度:", accuracy)

    accuracy = model.score(dataset["test_input"], dataset["test_label"])
    print("模型在测试数据集上的精度:", accuracy)


def kan_regression(dataset_fn, epoches=20):
    dataset = dataset_fn(extend=True)
    model   = KAN(width=[2, 2, 1], grid=3, k=3)

    def train_acc():
        return torch.mean((torch.round(model(dataset['train_input'])[:,0]) == dataset['train_label'][:,0]).float())

    def test_acc():
        return torch.mean((torch.round(model(dataset['test_input'])[:,0]) == dataset['test_label'][:,0]).float())

    results = model.train(dataset, opt="LBFGS", steps=epoches, metrics=(train_acc, test_acc))
    # print(results['train_acc'][-1], results['test_acc'][-1])
    print(results['train_acc'][-5:], results['test_acc'][-5:])


def kan_classify(dataset_fn, epoches=20):
    dataset = dataset_fn(extend=False)
    model = KAN(width=[2, 2], grid=3, k=3)

    def train_acc():
        return torch.mean((torch.argmax(model(dataset['train_input']), dim=1) == dataset['train_label']).float())

    def test_acc():
        return torch.mean((torch.argmax(model(dataset['test_input']), dim=1)  == dataset['test_label']).float())

    results = model.train(dataset, opt="LBFGS", steps=epoches, 
                        metrics=(train_acc, test_acc), 
                        loss_fn=torch.nn.CrossEntropyLoss())
    # print(results['train_acc'][-1], results['test_acc'][-1])
    print(results['train_acc'][-5:], results['test_acc'][-5:]) 


if __name__ == "__main__":
    # print(dataset['test_label'])
    # 模型在训练数据集上的精度: 0.995
    # 模型在测试数据集上的精度: 0.991
    # random_forest(dataset_fn=create_dataset, n_estimators=3, max_depth=6)

    # kan_regression(dataset_fn=create_dataset, epoches=20)

    kan_classify(dataset_fn=create_dataset, epoches=30)